import os
from ExtendScripts.DicomAutoConvert.DicomQuality.predict_class import *

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
weights_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "DicomQuality/resNet34.pth")

model = load_model(device, weights_path, num_classes=5)

class DicomQualityAdapter:
    @staticmethod
    def convert_nrrd_to_png(nrrd_path: str, png_path: str):
        process_single_file(nrrd_path, png_path)

    @staticmethod
    def read_classify_result(img_path: str):
        transform = transforms.Compose([
            transforms.ToTensor()
        ])
        img = preprocess_image(img_path, transform).to(device)
        with torch.no_grad():
            output = torch.squeeze(model(img)).cpu()
            predict = torch.softmax(output, dim=0)
            predict_cla = torch.argmax(predict).item()
            return int(predict_cla)

